#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   web_socket_item.py
@Time    :   2024-07-23 19:16:22
@Author  :   chakcy 
@Email   :   947105045@qq.com
@description   :   websocket 处理类
'''

import json
import logging
import asyncio
from abc import ABC, abstractmethod
from typing import Any, Protocol

from fastapi import WebSocket, Request

from .check_user_login_status import checkUserLoginStatus
from ..dependencies.jwt import decode_token
from ..dependencies.redis import RedisDependency

logging.basicConfig(level=logging.INFO)


class WebSocketTask(ABC):
    """
    WebSocketTask 基类

    用于处理 websocket 连接，接收数据，处理业务逻辑，返回结果，并返回给客户端

    属性:
    - handle: Callable[[Dict], Tuple[Any, Callable[[Any, WebSocket], Any]], 处理 websocket 连接的函数

    方法:
    - business: Callable[[Dict], Any], 处理业务逻辑的函数
    - callback: Callable[[Any, WebSocket], Any], 处理结果的回调函数

    注意：
    - 子类必须实现 business 方法，用于处理业务逻辑
    - 子类可以实现 callback 方法，用于处理结果的回调
    """
    def handle(self, data):
        return self.business(data), self.callback

    @abstractmethod
    def business(self, data):
        pass

    async def callback(self, result, websocket: WebSocket):
        pass

class WebSocketFunc(Protocol):
    async def __call__(self, websocket: WebSocket, 
                       redis: RedisDependency, 
                       *args, **kwargs) -> Any:
        ...

def websocket(func: WebSocketFunc) -> WebSocketFunc:
    """
    websocket 装饰器
    必须提供 websocket 和 redis 两个参数

    参数:
    - func: WebSocketFunc, 被装饰的函数
    - websocket: WebSocket, websocket 对象
    - redis: RedisDependency, redis 依赖对象

    返回:
    - wrapper: WebSocketFunc, 装饰后的函数

    注意：函数传参时，必须使用关键字参数，且必须提供 websocket 和 redis 两个参数
    """
    async def wrapper(*args, **kwargs):
        websocket:WebSocket = kwargs.get("websocket") # type: ignore
        redis:RedisDependency = kwargs.get("redis") # type: ignore

        if not websocket or not redis:
            raise ValueError("websocket and redis must be provided")
        
        try:
            token = websocket.query_params.get("token")
            if not token:
                await websocket.close(code=4001, reason="Token is missing in query params")
                return
            
            token = f"Bearer {token}"
            data = decode_token(token)
            if data is None:
                await websocket.close(code=4002, reason="Invalid token")
                return
        
            request = Request(scope={"type": "websocket", "headers": []})
            request.state.data = data
            request.state.token = token

            if await checkUserLoginStatus(request, redis) != 0:
                await websocket.close(code=4003, reason="User is not logged in")
                return

            await websocket.accept()
            while True:
                try:
                    data = await asyncio.wait_for(websocket.receive_text(), timeout=60*10)
                    logging.info(f"Received data: {data}")
                except asyncio.TimeoutError:
                    await websocket.close(code=4004, reason="Timeout")
                    break
                except RuntimeError:
                    break

                try:
                    json_data = json.loads(data)
                except json.JSONDecodeError:
                    json_data = data

                try:
                    business = await func(*args, **kwargs)
                    result, callback = business(json_data)
                except Exception as e:
                    logging.error(f"Error processing business logic: {e}")
                    await websocket.send_text(json.dumps({"error": str(e)}))

                try:
                    if callback:
                        asyncio.create_task(callback(result, websocket))
                except Exception as e:
                    logging.error(f"Error processing callback: {e}")
                    await websocket.send_text(json.dumps({"error": str(e)}))

        except Exception as e:
            try:
                await websocket.close(code=4000, reason=str(e))
            except RuntimeError:
                pass
            
    return wrapper